#!/usr/bin/python

'''
addClusterLabels.py:
	input: edge list and a set of clusters   
	output: edge list with clusters prepended to vertex name

usage:
	addClusterLabels.py edgelist clusters  > output
		all arguments are filenames

Created on May 1, 2011
@author: ncrowell
'''

import getopt
import sys

# creates a dictionary to map node names to the cluster they belong to.
# for now assumes one cluster per line in input file:
# 	clustername: vertex1 vertex2 vertex3 ...
def readClusters(f):
	clustermap = dict()
	for line in f.readlines():
		if (line.strip().startswith('#')):
			continue
		s = line.replace(':',' ').split()
		if ( len(s) > 0 ):
			clustername = s[0]
			for v in s[1:]:
				clustermap[v] = clustername
	return clustermap

def convertEdgeList(fin, clustermap, fout):
	for line in fin.readlines():
		if (line.strip().startswith('#')):
			continue
		# u v[ w]
		s = line.split()
		if ( len(s) > 0 ):
			u = s[0]
			v = s[1]
			w = ""
			if ( len(s) == 3):
				w = s[2]
			u = clustermap[u] + "/" + u		
			v = clustermap[v] + "/" + v
			fout.write(u + " " + v + " " + w + "\n");

def main():
	if ( len(sys.argv) != 3 ): 
		print ("usage:\n" +
			"\taddClusterLabels.py edgelist clusters  > output\n" +
			"\t\tall arguments are filenames\n")
		sys.exit()	

	inputfile = file(sys.argv[1], 'r')
	clusterfile = file(sys.argv[2], 'r')
	outputfile = sys.stdout

	# process cluster file
	clustermap = readClusters(clusterfile)
	clusterfile.close()

	# generate edge list with clusters prepended
	convertEdgeList(inputfile, clustermap, outputfile)

	inputfile.close()

if __name__ == '__main__':
	main()
